import os
import json
import logging
from typing import Dict, List, Any, Optional

logger = logging.getLogger(__name__)

class AnnotationCache:
    """
    可视化标注缓存管理器
    负责高效地加载、缓存和验证标注数据
    """
    
    def __init__(self, cache_dir: Optional[str] = None):
        """
        初始化标注缓存
        
        Args:
            cache_dir: 缓存存储目录，如果是None则仅使用内存缓存
        """
        # 内存缓存: {video_id: {frame_number: [annotations]}}
        self.memory_cache = {}
        self.cache_dir = cache_dir
        self.enabled = True
        
    def add_annotations(self, video_id: str, frame: int, annotations: List[Dict[str, Any]]):
        """
        添加标注到缓存
        
        Args:
            video_id: 视频ID
            frame: 帧号
            annotations: 标注列表
        """
        if not self.enabled:
            return
            
        if video_id not in self.memory_cache:
            self.memory_cache[video_id] = {}
            
        # 创建标注的深拷贝，避免对原始数据的修改
        clean_annotations = []
        for ann in annotations:
            clean_ann = {
                "frame": ann.get("frame"),
                "bbox": ann.get("bbox"),
                "color_id": ann.get("color_id"),
                "modality": ann.get("modality")
            }
            clean_annotations.append(clean_ann)
            
        self.memory_cache[video_id][frame] = clean_annotations
        
        # 如果配置了磁盘缓存，保存到磁盘
        if self.cache_dir:
            self._save_to_disk(video_id, frame, clean_annotations)
            
    def get_annotations(self, video_id: str, frame: int) -> Optional[List[Dict[str, Any]]]:
        """
        从缓存获取标注
        
        Args:
            video_id: 视频ID
            frame: 帧号
            
        Returns:
            标注列表或None（如果缓存中不存在）
        """
        if not self.enabled:
            return None
            
        # 首先尝试从内存缓存获取
        if video_id in self.memory_cache and frame in self.memory_cache[video_id]:
            # 返回深拷贝，避免修改缓存
            return [ann.copy() for ann in self.memory_cache[video_id][frame]]
            
        # 如果配置了磁盘缓存，尝试从磁盘加载
        if self.cache_dir:
            return self._load_from_disk(video_id, frame)
            
        return None
        
    def clear(self, video_id: Optional[str] = None, frame: Optional[int] = None):
        """
        清除缓存
        
        Args:
            video_id: 要清除的视频ID，如果为None则清除所有视频
            frame: 要清除的帧号，如果为None则清除指定视频的所有帧
        """
        if video_id is None:
            # 清除所有缓存
            self.memory_cache = {}
            if self.cache_dir:
                self._clear_disk_cache()
        elif frame is None:
            # 清除指定视频的所有缓存
            if video_id in self.memory_cache:
                del self.memory_cache[video_id]
            if self.cache_dir:
                self._clear_disk_cache(video_id)
        else:
            # 清除指定视频的指定帧缓存
            if video_id in self.memory_cache and frame in self.memory_cache[video_id]:
                del self.memory_cache[video_id][frame]
            if self.cache_dir:
                self._clear_disk_cache(video_id, frame)
                
    def _save_to_disk(self, video_id: str, frame: int, annotations: List[Dict[str, Any]]):
        """保存标注到磁盘缓存"""
        if not self.cache_dir:
            return
            
        try:
            video_dir = os.path.join(self.cache_dir, video_id)
            os.makedirs(video_dir, exist_ok=True)
            
            cache_file = os.path.join(video_dir, f"frame_{frame}.json")
            with open(cache_file, 'w', encoding='utf-8') as f:
                json.dump(annotations, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存标注到磁盘缓存失败: {e}")
            
    def _load_from_disk(self, video_id: str, frame: int) -> Optional[List[Dict[str, Any]]]:
        """从磁盘缓存加载标注"""
        if not self.cache_dir:
            return None
            
        try:
            cache_file = os.path.join(self.cache_dir, video_id, f"frame_{frame}.json")
            if not os.path.exists(cache_file):
                return None
                
            with open(cache_file, 'r', encoding='utf-8') as f:
                annotations = json.load(f)
                
            # 添加到内存缓存
            if video_id not in self.memory_cache:
                self.memory_cache[video_id] = {}
            self.memory_cache[video_id][frame] = annotations
            
            # 返回深拷贝
            return [ann.copy() for ann in annotations]
        except Exception as e:
            logger.error(f"从磁盘缓存加载标注失败: {e}")
            return None
            
    def _clear_disk_cache(self, video_id: Optional[str] = None, frame: Optional[int] = None):
        """清除磁盘缓存"""
        if not self.cache_dir or not os.path.exists(self.cache_dir):
            return
            
        try:
            if video_id is None:
                # 清除所有视频缓存
                for vdir in os.listdir(self.cache_dir):
                    vdir_path = os.path.join(self.cache_dir, vdir)
                    if os.path.isdir(vdir_path):
                        for f in os.listdir(vdir_path):
                            os.remove(os.path.join(vdir_path, f))
                        os.rmdir(vdir_path)
            elif frame is None:
                # 清除指定视频的所有缓存
                vdir_path = os.path.join(self.cache_dir, video_id)
                if os.path.exists(vdir_path):
                    for f in os.listdir(vdir_path):
                        os.remove(os.path.join(vdir_path, f))
                    os.rmdir(vdir_path)
            else:
                # 清除指定视频的指定帧缓存
                cache_file = os.path.join(self.cache_dir, video_id, f"frame_{frame}.json")
                if os.path.exists(cache_file):
                    os.remove(cache_file)
        except Exception as e:
            logger.error(f"清除磁盘缓存失败: {e}")
